#include <opencv2/opencv.hpp>
#include <iostream>
#include <fstream>
#include <stdio.h>
using namespace cv;
using namespace std;


Mat g_srcImage, g_dstImage;
Mat dst,last_dst;

int g_nFillMode = 1;

int g_nLowDifference = 20, g_nUpDifference = 20;

int g_nConnectivity = 4;

int g_nNewMaskVal = 255;
int labelClass = 0; 
int image_count = 224;
Point point1, point2;

 
void imageClear(Mat& img)  //clean image to make sure 
{
	for (int i = 0; i < img.rows; i++)
	{
		uchar* data = img.ptr<uchar>(i);
		for (int j = 0; j < img.cols; j++)
		{
			if ((data[3 * j] == 255 && data[3 * j + 1] == 255 && data[3 * j + 2] == 255) ||
				(data[3 * j] == 123 && data[3 * j + 1] == 123 && data[3 * j + 2] == 123)){
				data[3 * j] = 0; data[3 * j + 1] = 0;  data[3 * j + 2] = 0;
			}

		}
	}
}


static void onMouse(int event, int x, int y, int, void*)
{
 	if (event != CV_EVENT_LBUTTONDOWN)
		return;

	Point seed = Point(x, y);

	int LowDifference = g_nFillMode == 0 ? 0 : g_nLowDifference;
	int UpDifference = g_nFillMode == 0 ? 0 : g_nUpDifference; 
	int flags = g_nConnectivity + (g_nNewMaskVal << 8) +
		(g_nFillMode == 1 ? CV_FLOODFILL_FIXED_RANGE : 0); 

	Rect ccomp;

	Scalar label_0(255, 255, 255);
	Scalar label_1(255, 0, 0);  //label1  blue
	Scalar label_2(0, 255, 0);  //label1  green
	Scalar label_3(0, 0, 255);
	Scalar label_4(0, 125, 125);
	

	dst = g_dstImage; 
	g_dstImage.copyTo(last_dst);
	int area;

	if (dst.ptr<uchar>(y)[3 * x] == 0 && dst.ptr<uchar>(y)[3 * x + 1 ]== 0 && dst.ptr<uchar>(y)[3 * x + 2] == 0)
	{
		area = floodFill(dst, seed, label_0, &ccomp, Scalar(LowDifference, LowDifference, LowDifference),
			Scalar(UpDifference, UpDifference, UpDifference), flags);
	}

	if (labelClass == 0)
		area = floodFill(dst, seed, label_1, &ccomp, Scalar(LowDifference, LowDifference, LowDifference),
		Scalar(UpDifference, UpDifference, UpDifference), flags);
	else if (labelClass == 1)
		area = floodFill(dst,  seed, label_2, &ccomp, Scalar(LowDifference, LowDifference, LowDifference),
		Scalar(UpDifference, UpDifference, UpDifference), flags);
	else if (labelClass == 2)
		area = floodFill(dst, seed, label_3, &ccomp, Scalar(LowDifference, LowDifference, LowDifference),
		Scalar(UpDifference, UpDifference, UpDifference), flags);
	else if (labelClass == 3)
		area = floodFill(dst,  seed, label_4, &ccomp, Scalar(LowDifference, LowDifference, LowDifference),
		Scalar(UpDifference, UpDifference, UpDifference), flags);

	imshow("image", dst);
	cout << area << "pixels are filled\n";
}

static void draw_line(int event, int x, int y, int, void*){

	

	if (event != CV_EVENT_LBUTTONDOWN&&event != CV_EVENT_LBUTTONUP)
		return;

	if (event == CV_EVENT_LBUTTONDOWN){
		point1 = Point(x, y);
	}

	if (event == CV_EVENT_LBUTTONUP){
		dst = g_dstImage;
		g_dstImage.copyTo(last_dst);
		point2 = Point(x, y);
		line(dst, point1, point2, Scalar(123, 123, 123));
		imshow("image", dst);
	}	
}

static void draw_rec(int event, int x, int y, int, void*){



	if (event != CV_EVENT_LBUTTONDOWN&&event != CV_EVENT_LBUTTONUP)
		return;

	if (event == CV_EVENT_LBUTTONDOWN){
		point1 = Point(x, y);
	}

	if (event == CV_EVENT_LBUTTONUP){
		dst = g_dstImage;
		g_dstImage.copyTo(last_dst);
		point2 = Point(x, y);
		double R = sqrt((point2.x - point1.x)*(point2.x - point1.x) + (point1.y - point2.y)*(point1.y - point2.y));
		double c_x = (point2.x + point1.x) / 2;
		double c_y = (point1.y + point2.y) / 2;
		for (int i = int(c_y - R); i <int(c_y + R); i++)
		{
			uchar* data = dst.ptr<uchar>(i);
			for (int j = int(c_x - R); j <int(c_x + R); j++)
			{
				if (((i - c_y)*(i - c_y) + (j - c_x)*(j - c_x)) <= R*R)
				{
					data[3 * j] = 255; data[3 * j + 1] = 255;  data[3 * j + 2] = 255;
				}
			}

			imshow("image", dst);
		}
	}
}

void dataAug(int start, int end)
{
	for(int i = start; i <= end; i++)
	{
		stringstream ss; ss << i;
		cout << "SampleImages/image" + ss.str() + ".png";
		cv::Mat sample = cv::imread("SampleImages/image" + ss.str() + ".png");
		cv::Mat gt = cv::imread("GT_labelled/" + ss.str() + ".png");
		cv::Mat sampele_flipped,gt_flipped;

		cv::flip(sample,sampele_flipped,1);
		cv::flip(gt,gt_flipped,1);
		// cv::imshow("s",sampele_flipped);
		// cv::imshow("gt",gt_flipped);
		// cv::waitKey();

		stringstream ss_; ss_ << i + 306 + 1;
		cv::imwrite("SampleImages/image" + ss_.str() + ".png",sampele_flipped);
		cv::imwrite("GT_labelled/" + ss_.str() + ".png",gt_flipped);
		cout << "image" << ss_.str() << " saved!" << endl;
	}
}

void loadImage(std::string data_file,std::vector<cv::Mat>& res)
{
    res.clear();
    std::ifstream file_image;
    file_image.open(data_file.c_str());
    while(!file_image.eof())
    {
        std::string s;
        getline(file_image,s);
        if(!s.empty())
        {
          cout << s << endl;
          cv::Mat temp = cv::imread(s);
          // cv::imshow("img",temp);
          // cv::waitKey();
          res.push_back(temp);
        }
    }
}

int main(int argc, char** argv)
{
	if(argc != 3){
		cout << "error! IMAGE_NAME_FILE IMAGE_COUNT" << endl;
		return -1;
	}

	std::vector<cv::Mat> mask_vec;
	loadImage(argv[1],mask_vec);
	int image_count = strtol(argv[2], NULL, 10);  //int arg = strtol(argv[1], NULL, 10);
	int image_cnt_index = image_count % 500;
	while (true)
	{
		stringstream ss_; ss_ << image_count;

//------------------------------------------------------------------------------
//skip the first image of each turn.Mask is corrupted
		if((image_cnt_index % 10) == 0){
			image_cnt_index++;
			image_count++;
			continue;
		}
//-------------------------------------------------------------------------------


		g_srcImage = mask_vec[image_cnt_index];//imread("sub_image/image" + ss_.str() + ".png", 1);
		cout << "current sample " << image_count << endl;
		labelClass = 0;

		if (!g_srcImage.data) { cout << "image read error!"; return false; }

		g_srcImage.copyTo(g_dstImage);

		namedWindow("image", CV_WINDOW_AUTOSIZE);

		createTrackbar("maximum negative", "image", &g_nLowDifference, 255, 0);
		createTrackbar("maximum positive", "image", &g_nUpDifference, 255, 0);		
		setMouseCallback("image", onMouse, 0);
		while (1)
		{
			imshow("image", g_dstImage );
			stringstream ss;
			int c = waitKey(0);
			//ESC quit
			if ((c & 255) == 27)
			{
				cout << "quit...........\n";
				break;
			}
			
			switch ((char)c)
			{
			case '1':
				labelClass =(labelClass+1)%4;
				cout << "switch to class" << labelClass << endl;
				break;
			case '2':
				ss << image_count++;
				image_cnt_index++;
				imageClear(dst);
				imwrite("GT_labelled/" + ss.str() + ".png", dst);
				cout << "label saved!\n";

				//cout << "Please input index for the first image\n";
				//cin >> image_count;
				break;
			case '3':
				cout << "move back to original image\n";
				g_srcImage.copyTo(g_dstImage);
				labelClass = 0;
				break;
			case '5':
				cout << "flood fill mode 1\n";
				g_nFillMode = 1;
				break;
			case '6':
				cout << "flood fill mode 2\n";
				g_nFillMode = 2;
				break;
			case '7':
				cout << "g_nConnectivity 4\n";
				g_nConnectivity = 4;
				break;
			case '8':
				cout << "g_nConnectivity 8\n";
				g_nConnectivity = 8;
				break;
			case '9':
				cout << "flood fill mode 0\n";
				g_nFillMode = 0;
				break;
			case'd':
				cout << "Now please draw the segmenting line" << endl;
				setMouseCallback("image", draw_line, 0);
				break;
			case'c':
				cout << "Now please color the map" << endl;
				setMouseCallback("image", onMouse, 0);
				break;
			case 'r':
				cout << "Your selection will be cleared" << endl;
				setMouseCallback("image", draw_rec, 0);
				break;
			case'z':
				cout << "one step back" << endl;
				last_dst.copyTo(g_dstImage);
				imshow("image", g_dstImage);
				break;
			case'n':
				cout << "Please input the number you want" << endl;
				cin >> image_count;
				break;			
			case 'q':
				cout << "quit!\n";
				return 0;	
				}
		}
	
	}
	return 0;
}
